# Easy data augmentation techniques for text classification
# Jason Wei and Kai Zou
# Fixed by GMFTBY, 2019.8.3

import jieba
from eda import eda
import argparse
from tqdm import tqdm


# generate more data with standard augmentation
def gen_eda(input_file, output_file, alpha, num_aug=9):
    print('Begin to generate the augmentation data, num_aug:{}, alpha:{}'.format(num_aug, alpha))

    # load origin data
    with open(input_file, 'r', encoding='utf-8') as f:
        lines = [' '.join(jieba.cut(l)) for l in f.readlines()]

    total = len(lines)
    print(r'Total lines: {}'.format(total))
    with open(output_file, 'w', encoding='utf-8') as f:
        pbar = tqdm(enumerate(lines))
        for i, line in pbar:
            sentence = line.strip()
            # 保留原句子
            f.write(sentence.replace(' ', '') + '\n')
            aug_sentences = eda(sentence,
                                alpha_sr=alpha,  # Synonym Replacement (SR)
                                alpha_ri=alpha,  # Random Insertion (RI)
                                alpha_rs=alpha,  # Random Swap (RS)
                                p_rd=alpha,  # Random Deletion (RD)
                                num_aug=num_aug)
            for aug_sentence in aug_sentences:
                # 相同的不保留
                if aug_sentence != sentence.replace(' ', ''):
                    f.write(aug_sentence + '\n')
            pbar.set_description(r'Progress: {} / {}, {}'.format(i, total, 100 * round(i / total, 4)))

        pbar.close()
    print(r"Over. EDA-Chinese for {} to {} with num_aug={}".format(input_file, output_file, num_aug))


if __name__ == "__main__":
    # args parse
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", required=True, type=str,
                    help="input file of un-augmented data")
    ap.add_argument("--output", required=False, type=str,
                    help="output file of augmented data")
    ap.add_argument("--num_aug", required=False, type=int, default=4,
                    help="number of augmented sentences per original sentence")
    ap.add_argument("--alpha", required=False, type=float, default=0.1,
                    help="percent of words in each sentence to be changed")
    args = ap.parse_args()

    output = None
    if args.output:
        output = args.output
    else:
        from os.path import dirname, basename, join

        output = join(dirname(args.input), 'eda_' + basename(args.input))

    # number of augmented sentences to generate per original sentence
    num_aug = args.num_aug
    alpha = args.alpha

    # generate augmented sentences and output into a new file
    gen_eda(args.input, output, alpha=alpha, num_aug=num_aug)
